/*
 * Copyright 2013-2018 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.hundsun.aitest.controller;

import com.alibaba.fastjson.JSONArray;
import com.hundsun.aitest.common.BladePlatformConfig;
import com.hundsun.aitest.common.Result;
import com.hundsun.aitest.model.ai.AutoCase;
import com.hundsun.aitest.model.ai.ChatRequest;
import com.hundsun.aitest.model.ai.Message;
import com.hundsun.aitest.model.blade.TestCaseRequest;
import com.hundsun.aitest.model.plat.InterfaceDict;
import com.hundsun.aitest.model.plat.TableInformation;
import com.hundsun.aitest.model.testanalysistree.LeafNode;
import com.hundsun.aitest.service.AutoCaseParser;
import com.hundsun.aitest.service.LargeModelService;
import com.hundsun.aitest.service.TestCasesService;
import com.hundsun.aitest.service.TestInterfaceService;
import com.hundsun.aitest.util.IndentationParser;
import com.hundsun.aitest.util.TaskProgress;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.ResponseBody;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

import static com.hundsun.aitest.common.CaseTypeEnum.AUTO;

@Controller
@Slf4j
public class AIGenerateController {

    @Autowired
    private LargeModelService largeModelService;

    @Autowired
    private TestCasesService testCasesService;

    @Autowired
    private AutoCaseParser<AutoCase> testCaseParser;

    @Autowired
    private BladePlatformConfig bladePlatformConfig;

    @Autowired
    private TestInterfaceService interfaceService;

    // 任务进度存储
    private final Map<String, TaskProgress> taskProgressMap;

    @Autowired
    public AIGenerateController(ConcurrentHashMap<String, TaskProgress> taskProgressMap) {
        this.taskProgressMap = taskProgressMap;
    }

    @PostMapping("/aigenerate")
    @ResponseBody
    public Result<String> handlePostRequest(@RequestBody Map<String, String> requestBody) {
        Result<String> result = new Result<>();
        result.setSuccess(true);
        result.setData("生成中...");

        String content = requestBody.get("content");
        String account = requestBody.getOrDefault("account", "");
        String nodePath = requestBody.get("nodePath");
        String projectName = requestBody.get("projectName");

        String apiToken = getApiToken(projectName);
        account = StringUtils.defaultIfEmpty(account, getAccount(projectName));

        String taskId = UUID.randomUUID().toString();
        // 初始化任务进度
        taskProgressMap.put(taskId, new TaskProgress(1));

        processNodeAsync(taskId, apiToken, account, nodePath, projectName, "", content, 1);

        result.setData(taskId);
        return result;
    }

    @PostMapping("/batch_aigenerate")
    @ResponseBody
    public Result<String> handleBatchPostRequest(@RequestBody Map<String, String> requestBody) {
        Result<String> result = new Result<>();
        result.setSuccess(true);
        result.setData("生成中...");

        String content = requestBody.get("content");
        String account = requestBody.getOrDefault("account", "");
        String nodePath = requestBody.get("nodePath");
        String projectName = requestBody.get("projectName");
        String prepareData = requestBody.getOrDefault("prepareData", "");

        String apiToken = getApiToken(projectName);
        account = StringUtils.defaultIfEmpty(account, getAccount(projectName));

        List<LeafNode> analysisList = IndentationParser.parse(content);
        String taskId = UUID.randomUUID().toString();

        // 初始化任务进度
        taskProgressMap.put(taskId, new TaskProgress(analysisList.size()));

        processNodesSequentially(taskId, apiToken, account, nodePath, projectName, analysisList);

        result.setData(taskId);
        return result;
    }

    private void processNodesSequentially(String taskId, String apiToken, String account, String nodePath, String projectName, List<LeafNode> nodes) {
        AtomicInteger index = new AtomicInteger(0);
        List<CompletableFuture<Void>> futures = nodes.stream()
                .map(node -> processNodeAsync(taskId, apiToken, account, nodePath, projectName, node.getPath(), node.getValue(), index.incrementAndGet()))
                .collect(Collectors.toList());

        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
                .thenRun(() -> {
                    TaskProgress progress = taskProgressMap.get(taskId);
                    if (progress.getCompleted() == progress.getTotal()) {
                        updateTaskProgress(taskId, "任务处理完成，所有节点处理成功");
                    } else {
                        updateTaskProgress(taskId, "任务处理完成，部分节点处理失败");
                    }
                })
                .exceptionally(t -> {
                    log.error("处理节点时发生错误", t);
                    updateTaskProgress(taskId, "任务处理完成，部分节点处理失败");
                    return null;
                });
    }


    private CompletableFuture<Void> processNodeAsync(String taskId, String apiToken, String account, String nodePath,
                                                     String projectName, String path, String value, int index) {
        return CompletableFuture.runAsync(() -> {
            try {
                processNode(apiToken, account, nodePath, projectName, path, value);
                updateTaskProgress(taskId, "成功处理节点 " + index + ": " + path);
            } catch (Exception e) {
                log.error("异步处理节点时发生错误，节点路径：{}", path, e);
                updateTaskProgress(taskId, "错误处理节点 " + index + ": " + path);
                throw new RuntimeException(e);
            }
        });
    }

    private synchronized void updateTaskProgress(String taskId, String message) {
        TaskProgress progress = taskProgressMap.get(taskId);
        if (progress != null) {
            progress.incrementCompleted();
            int percent = (int) ((double) progress.getCompleted() / progress.getTotal() * 100);
            log.info("任务 {} 进度更新: {}% - {}", taskId, percent, message);
            progress.setMessage(message);
        }
    }

    /*@GetMapping("/taskStatus")
    @ResponseBody
    public Result<String> getTaskStatus(@RequestParam String taskId) {
        Result<String> result = new Result<>();
        try {
            TaskProgress progress = taskProgressMap.get(taskId);
            if (progress != null) {
                result.setSuccess(true);
                int percent = (int) ((double) progress.getCompleted() / progress.getTotal() * 100);
                result.setData(String.valueOf(percent));
            } else {
                result.setSuccess(false);
                result.setData("任务ID未找到");
            }
        } catch (Exception e) {
            log.error("查询任务状态时发生错误，任务ID：{}", taskId, e);
            result.setSuccess(false);
            result.setData("任务状态查询失败");
        }
        return result;
    }
*/
    private String processNode(String apiToken, String account, String nodePath, String projectName, String path,
                               String value) throws Exception {
        String currentNodePath = nodePath + (path.isEmpty() ? "" : "/" + path);
        InterfaceDict interfaceResult = testCasesService.patternAndGetInterfaceDetail(projectName,
                path + (value.isEmpty() ? "" : "/" + value));
        String interfaceDesc = buildInterfaceDescription(interfaceResult);
        String tableStructures = buildTableStructureDescription(interfaceResult);
        // 构建prompt
        String prompt = largeModelService.buildPrompt(path + (value.isEmpty() ? "" : "/" + value),
                interfaceDesc, tableStructures);
        ChatRequest request = createChatRequest(prompt);

        String response = largeModelService.sendRequest(request);
        log.info("AI返回结果：\n{},\nrequest:\n{}", response, request);

        String jsonString = largeModelService.parseResponse(response);
        AutoCase autoCase = testCaseParser.parse(jsonString, AUTO.getValue(), interfaceResult.getScriptId());

        TestCaseRequest testCaseRequest = testCasesService.buildTestCaseRequest(apiToken, account, currentNodePath,
                autoCase, interfaceResult.getScriptId());
        String uploadResult = testCasesService.uploadTestCase(testCaseRequest);

        return uploadResult;
    }

    private String getApiToken(String projectName) {
        BladePlatformConfig.ProjectConfig projectResult = bladePlatformConfig.getBladePlatform().get(projectName);
        return projectResult != null ? projectResult.getApiKey() : "";
    }

    private String getAccount(String projectName) {
        BladePlatformConfig.ProjectConfig projectResult = bladePlatformConfig.getBladePlatform().get(projectName);
        return projectResult != null ? projectResult.getAccount() : "";
    }

    private String buildInterfaceDescription(InterfaceDict interfaceResult) {
        String interfaceDesc = "    URL：" + interfaceResult.getUrl()
                + "\n        入参：" + interfaceResult.getParamsList()
                + "\n        出参：" + interfaceResult.getResponseParamsList();
        if(interfaceResult.getCommonParamsDict()!=null){
            interfaceDesc += "\n        公共字典项：" + interfaceResult.getCommonParamsDict();
        }
        log.debug("buildInterfaceDescription: {}", interfaceDesc);
        return interfaceDesc;

    }

    private String buildTableStructureDescription(InterfaceDict interfaceResult){
        List<TableInformation> tableStructures = interfaceResult.getTableStructures();
        String tableStructureDesc = "    ";
        if(tableStructures !=null && !tableStructures.isEmpty()) {
            for(TableInformation tableStructure : tableStructures){
                tableStructureDesc += tableStructure.getTableName() + ":" + JSONArray.toJSONString(tableStructure.getTableStructure()) + "\n";
                // tableStructureDesc += "[" + tableStructure.getTableName() + "]的表数据示例:" + tableStructure.getDataExample() + "\n";
            }
        }
        return tableStructureDesc;
    }

    private ChatRequest createChatRequest(String prompt) {
        ChatRequest request = new ChatRequest();
        request.setMessages(Arrays.asList(new Message("user", prompt)));
        return request;
    }

}